package lib

func init() {
	Probs = append(Probs, Problem{
		Num:         347,
		Discription: "前k个高频元素",
		Level:       2,
		Labels: map[string]int{
			"大小堆": 1,
			"哈希表": 1,
			"多元排序": 1,
		},
	})
}

func TopKFrequent(nums []int, k int) []int {
	valCount := make(map[int]int, 0)
	for _, val := range nums {
		valCount[val]++
	}

	numCounts := make([][2]int, 0)
	for val, count := range valCount {
		numCounts = append(numCounts, [2]int{val, count})
	}

	buildHeap(numCounts)
	end := len(numCounts) - 1
	for i := 0; i < k; i++ {
		numCounts[0], numCounts[end] = numCounts[end], numCounts[0]
		end--
		adjustHeap(numCounts, 0, end)
	}

	res := make([]int, 0)
	for i := end + 1; i <= len(numCounts)-1; i++ {
		res = append(res, numCounts[i][0])
	}

	return res
}

func buildHeap(nums [][2]int) {
	for i := len(nums)/2 - 1; i >= 0; i-- {
		adjustHeap(nums, i, len(nums)-1)
	}
}

func adjustHeap(nums [][2]int, start int, end int) {
	topValue := nums[start]
	i := start
	for i*2+1 <= end {
		leftIndex := 2*i + 1
		rightIndex := leftIndex + 1
		biggerIndex := leftIndex
		if rightIndex <= end && nums[rightIndex][1] > nums[leftIndex][1] {
			biggerIndex = rightIndex
		}

		if topValue[1] >= nums[biggerIndex][1] {
			break
		}

		nums[i], nums[biggerIndex] = nums[biggerIndex], nums[i]
		i = biggerIndex
	}

	nums[i] = topValue
}